package mr

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"time"
)
import "log"
import "net/rpc"
import "hash/fnv"

// Map functions return a slice of KeyValue.
type KeyValue struct {
	Key   string
	Value string
}

type WorkerState string

const (
	WorkerIdle   WorkerState = "idle"
	WorkerBusy   WorkerState = "busy"
	WorkerFailed WorkerState = "failed"
)

type Worker struct {
	workerID           string // 唯一标识符
	coordinatorAddress string // 协调器的地址
	task               Task   // 当前正在进行的任务

	// 状态信息
	state                  WorkerState // "idle", "busy", "failed"
	lastTaskCompletionTime time.Time   // 最近一次任务完成的时间

	// 心跳机制相关
	heartbeatInterval time.Duration // 心跳间隔
	lastHeartbeatSent time.Time     // 上次发送心跳的时间

	// 网络连接
	connection net.Conn // 到协调器的连接
}

// NewWorker 初始化 Worker 实例
func NewWorker(coordinatorAddress string) *Worker {
	return &Worker{
		coordinatorAddress: coordinatorAddress,
		state:              WorkerIdle,
		heartbeatInterval:  10 * time.Second, // 默认的心跳间隔
	}
}

// *************** Worker 方法   ***************

// ConnectToCoordinator 用于初始化与协调器的连接
func (w *Worker) connectToCoordinator() error {
	conn, err := net.Dial("tcp", w.coordinatorAddress)
	if err != nil {
		return fmt.Errorf("failed to connect to coordinator at %s: %v", w.coordinatorAddress, err)
	}
	w.connection = conn
	// 连接时初始化为ip + 端口号
	w.workerID = conn.LocalAddr().String()
	log.Printf("成功连接协调器 at %s\n workerId为 %s\n", w.coordinatorAddress, w.workerID)
	return nil
}

// 定期发送心跳给协调器
func (w *Worker) sendHeartbeat() {
	for {
		if w.connection == nil {
			log.Println("No active connection to coordinator. Attempting to reconnect...")
			if err := w.connectToCoordinator(); err != nil {
				log.Printf("Failed to reconnect: %v", err)
				time.Sleep(w.heartbeatInterval)
				continue
			}
		}

		// 发送心跳逻辑
		_, err := w.connection.Write([]byte("HEARTBEAT"))
		if err != nil {
			log.Printf("Failed to send heartbeat: %v", err)
			w.connection.Close()
			w.connection = nil
			continue
		}

		log.Println("Heartbeat sent.")
		time.Sleep(w.heartbeatInterval)
	}
}
func (w *Worker) executeMapTask(mapf func(string, string) []KeyValue) {
	dir, _ := os.Getwd()
	filename := filepath.Join(filepath.Dir(dir), "in") + "\\" + w.task.FileName
	file, err := os.Open(filename)
	if err != nil {
		log.Fatalf("cannot open %v: %v", filename, err)
	}
	content, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatalf("cannot read %v: %v", filename, err)
	}
	kva := mapf(filename, string(content))

	// 4. 按照 key 对  reduce 任务的数量  进行 hash 平均分布reduce任务
	intermediate := make(map[int][]KeyValue)
	for _, kv := range kva {
		num := ihash(kv.Key) % w.task.ReduceNum
		intermediate[num] = append(intermediate[num], kv)
	}

	tempDir := filepath.Join(filepath.Dir(dir), "temp")
	//5.存储到json文件下
	for i, kvs := range intermediate {
		tf, err := os.CreateTemp(tempDir, "mr-temp")

		if err != nil {
			log.Fatalf("cannot create temp file: %v", err)
		}
		encoder := json.NewEncoder(tf)
		for _, kv := range kvs {
			encoder.Encode(kv)
		}
		tf.Close()
		defer os.Remove(tf.Name())
		dir, _ := os.Getwd()
		targetDir := filepath.Join(filepath.Dir(dir), "map-out")
		// 构造目标文件名
		fileName := fmt.Sprintf(targetDir+"\\mr-%v-%v", w.task.TaskID, i)
		os.Rename(tf.Name(), fileName)
	}
}

func (w *Worker) executeReduceTask(reduce func(string, []string) string) {
	id := w.task.TaskID
	num := w.task.ReduceNum
	intermediate := make(map[string][]string)

	for i := 0; i < num; i++ {
		dir, _ := os.Getwd()
		targetDir := filepath.Join(filepath.Dir(dir), "map-out")
		fileName := fmt.Sprintf(targetDir+"\\mr-%v-%v", id, i)

		file, err := os.Open(fileName)
		if err != nil {
			log.Fatalf("cannot open %v: %v", fileName, err)
		}
		dec := json.NewDecoder(file)
		for {
			var kv KeyValue
			if err := dec.Decode(&kv); err == io.EOF {
				break
			} else if err != nil {
				log.Fatalf("cannot decode %v: %v", fileName, err)
			}
			intermediate[kv.Key] = append(intermediate[kv.Key], kv.Value)
		}
		file.Close()
	}
	dir, _ := os.Getwd()
	targetDir := filepath.Join(filepath.Dir(dir), "reduce-out")
	// 3. 执行 reduce函数，并将结果存储到输出
	outputFile, _ := os.Create(fmt.Sprintf(targetDir+"//mr-out-%d", w.task.TaskID))
	for k, v := range intermediate {
		cnt := reduce(k, v)
		fmt.Fprintf(outputFile, "%v %v\n", k, cnt)
	}
	outputFile.Close()

}

// *************** Worker 方法   ***************

// *************** Rpc接口  ***************

// callGetTask  请求任务
func (w *Worker) callGetTask() error {
	if w.state == WorkerBusy {
		return fmt.Errorf("worker is busy, waiting for a free slot")
	}
	getTaskArgs := GetTaskArgs{w.workerID}
	var res GetTaskReply
	err := w.call("Coordinator.GetTask", &getTaskArgs, &res)
	if err != nil {
		return err
	}
	log.Printf("[DEBUG]调用GetTask方法成功,返回值为 %#v", res.Task)
	//初始化Worker
	w.task = res.Task
	w.state = WorkerBusy
	log.Printf("workerID: %v 获得任务，任务id为：%v,任务类型为：%v", w.workerID, res.Task.TaskID, res.Task.TaskType)
	return nil
}

func (w *Worker) callTaskDone() error {
	if w.state != WorkerBusy {
		return fmt.Errorf("worker is not busy, can't call TaskDone")
	}
	taskDoneArgs := CompleteTaskArgs{w.task.TaskID, w.task.TaskType}
	var res CompleteTaskReply
	err := w.call("Coordinator.CompleteTask", &taskDoneArgs, &res)
	if err != nil {
		return err
	}
	log.Printf("[DEBUG]调用TaskDone方法成功 worker%v  完成任务 %v", w.workerID, w.task.TaskID)
	w.task.Status = Completed
	w.state = WorkerIdle
	return nil
}

// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
func (w *Worker) call(rpcname string, args interface{}, reply interface{}) error {
	c, err := rpc.DialHTTP("tcp", w.coordinatorAddress)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return nil
	}

	fmt.Println(err)
	return err
}

// *************** Rpc接口  ***************

// use ihash(key) % NReduce to choose the reduce
// Task number for each KeyValue emitted by Map.
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

// StartWorker main/mrworker.go calls this function.
// 启动Worker进程
func StartWorker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {
	worker := NewWorker(coordinatorIp + ":" + coordinatorPort)

	// 1. 连接协调器
	for {
		if err := worker.connectToCoordinator(); err != nil {
			log.Printf("Failed to initialize connection: %v", err)
			log.Println("Retrying in 10 seconds...")
			time.Sleep(10 * time.Second)
			continue
		}
		log.Println("Successfully connected to coordinator.")
		break
	}
	//// 2. 开启心跳协程
	//go worker.sendHeartbeat()

	// 3. 主循环：请求任务 -> 执行任务 -> 报告任务完成
	for {
		err := worker.callGetTask()
		if err != nil {
			log.Printf("Failed to get Task:")
			time.Sleep(5 * time.Second) // 等待一段时间后重试
			continue
		}

		// 根据任务类型执行相应的操作
		switch worker.task.TaskType {
		case MapTask:
			time.Sleep(5 * time.Second)
			log.Printf("Executing Map Task: %+v\n", worker.task)
			worker.executeMapTask(Map)
		case ReduceTask:
			log.Printf("Executing Reduce Task: %+v\n", worker.task)
			worker.executeReduceTask(Reduce)
		default:
			log.Printf("没有任务，需要等待")
			time.Sleep(5 * time.Second)
		}

		//完成任务
		if err := worker.callTaskDone(); err != nil {
			log.Printf("Failed to report Task done: %v", err)
		}

		log.Println("")
	}
}
